feat: pyproj CRS extension — generic reproject(x, y, src_crs, dst_crs) UDF#1
Open
Mmoncadaisla wants to merge 4 commits into
Open
feat: pyproj CRS extension — generic reproject(x, y, src_crs, dst_crs) UDF#1Mmoncadaisla wants to merge 4 commits into
Mmoncadaisla wants to merge 4 commits into
Conversation
…) UDF Adds xarray_sql/proj.py, an optional pyproj extension that registers an ST_Transform-style scalar UDF. The CRS pair is part of the query (any spelling pyproj.CRS accepts, and it may vary per row via ordinary SQL expressions) instead of being baked in at UDF registration time, which is how the geospatial benchmarks previously hard-coded it. All PROJ work runs on a dedicated pool of Python-owned worker threads: constructing a Transformer on a DataFusion runtime thread segfaults inside PROJ, while identical concurrent work on Python threads is stable (pyproj 3.7 / PROJ 9.5). Each pool thread caches one transformer per CRS pair, and pyproj releases the GIL during transforms, so the UDF is parallel across partitions — the previous single-chunk/serial-UDF workaround in the benchmarks is no longer needed. XarrayContext auto-registers reproject() when pyproj is installed (pip install xarray-sql[proj]); proj.register() is the explicit hook for plain SessionContexts or custom names. Benchmarks 07 and 09 now use the extension instead of their local hard-coded UDFs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The y-chunking exists to force several DataFusion partitions so the reproject() UDF demonstrably runs in parallel; say so explicitly instead of hiding it behind a computed chunk size. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The kernel materialized the two broadcast CRS string columns with to_pylist() — two Python object allocations per row before PROJ did any work. At benchmark scale that dominated everything: a 606M-pixel aggregate-only reprojection took 496s while the same scan without the UDF took 4.5s. Establish CRS uniqueness with pyarrow.compute.unique (vectorized C) instead. A batch with one CRS pair — the overwhelmingly common case — becomes a single vectorized PROJ call with no per-row Python. The per-row grouping path remains for genuinely varying CRS (e.g. a CASE expression picking the UTM zone). Same 606M-pixel run after: 16.0s (37.9M px/s), a 31x speedup, with bit-identical results and flat (~2GB) streaming memory. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Owner
Author
Scale test: all 9.13 billion pixels of a real national rasterRan the UDF over every pixel coordinate of a production 675 MB GeoTIFF (ERI Argentina v3, 123,459 × 73,963 ≈ 9.13e9 cells, WGS84 → UTM 20S) as an aggregate-only streaming query (nothing materialized). On an M-series MacBook:
Findings:
Script: 🤖 Generated with Claude Code |
Aligns with the base + add-ons direction: geo is the umbrella extra for CRS reprojection today and polygon/vector support later, rather than being named after the first feature it happens to ship. The module stays proj.py — it does CRS reprojection specifically. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fork-local copy of the PR proposed upstream as xqlsystems#225 — kept here for review and iteration on my own fork.
Summary
xarray_sql/proj.py(pip install xarray-sql[proj]): anST_Transform-style scalar UDFreproject(x, y, src_crs, dst_crs)returning a{x, y}struct. CRS pair lives in the query (any pyproj CRS spelling, may vary per row) instead of being hard-coded at registration time.pyproj.Transformeron DataFusion's tokio runtime threads segfaults inside PROJ, while identical concurrent work on Python threads is stable (pyproj 3.7.2 / PROJ 9.5.1). Per-thread transformer caches; pyproj releases the GIL, so partitions transform in parallel.XarrayContextauto-registersreproject()when pyproj is installed (same pattern as thecftime()UDF);proj.register(ctx, name=...)for plainSessionContexts.inf→ NaN; invalid CRS fails the query loudly.tests/test_proj.pyadded and pyproj added to thetestextra.Design context — why the CRS is a query argument: DataFusion's type system can't yet attach metadata like a CRS to a column type (extension types: apache/datafusion#12644; spatial support overall: apache/datafusion#7859), so an "implicit CRS" would have to hide in UDF registration state — the hard-coding this PR removes. When extension types land and GeoArrow columns carry their CRS (e.g. datafusion-contrib/datafusion-geo), a geometry-typed
ST_Transform(geom, dst)can arrive alongsidereproject()rather than replace it.See the upstream PR for the full write-up and testing caveats.
🤖 Generated with Claude Code